#pragma once
#include <brpc/channel.h>
#include <string>
#include <vector>
#include <unordered_map>
#include <mutex>
#include "logger.hpp"

namespace mag {

// 封装单个服务的信道管理类
class ServiceChannel {
public:
    using ptr = std::shared_ptr<ServiceChannel>;
    using ChannelPtr = std::shared_ptr<brpc::Channel>;

    // 构造函数，初始化服务名称和轮询计数器
    ServiceChannel(const std::string &name) : _service_name(name), _index(0) {}

    // 服务上线了一个节点时调用，新增信道
    void append(const std::string &host) {
        auto channel = std::make_shared<brpc::Channel>();
        brpc::ChannelOptions options;
        options.connect_timeout_ms = -1; // 设置连接超时为无限等待
        options.timeout_ms = -1;        // 设置请求超时为无限等待
        options.max_retry = 3;          // 设置最大重试次数
        options.protocol = "baidu_std"; // 使用百度标准协议

        // 初始化信道
        int ret = channel->Init(host.c_str(), &options);
        if (ret == -1) {
            LOG_ERROR("初始化{}-{}信道失败!", _service_name, host);
            return;
        }

        // 加锁保护，更新信道集合和主机映射
        std::unique_lock<std::mutex> lock(_mutex);
        _hosts.insert(std::make_pair(host, channel));
        _channels.push_back(channel);
    }

    // 服务下线了一个节点时调用，移除信道
    void remove(const std::string &host) {
        std::unique_lock<std::mutex> lock(_mutex);
        auto it = _hosts.find(host);
        if (it == _hosts.end()) {
            LOG_WARN("{}-{}节点删除信道时，没有找到信道信息！", _service_name, host);
            return;
        }

        // 从信道集合中移除对应的信道
        for (auto vit = _channels.begin(); vit != _channels.end(); ++vit) {
            if (*vit == it->second) {
                _channels.erase(vit);
                break;
            }
        }

        // 从主机映射中移除
        _hosts.erase(it);
    }

    // 使用RR轮转策略，获取一个Channel用于发起RPC调用
    ChannelPtr choose() {
        std::unique_lock<std::mutex> lock(_mutex);
        if (_channels.size() == 0) {
            LOG_ERROR("当前没有能够提供 {} 服务的节点！", _service_name);
            return ChannelPtr();
        }

        // 通过轮询选择一个信道
        int32_t idx = _index++ % _channels.size();
        return _channels[idx];
    }

private:
    std::mutex _mutex;                               // 互斥锁，保护多线程环境下的成员变量
    int32_t _index;                                  // 当前轮转下标计数器
    std::string _service_name;                       // 服务名称
    std::vector<ChannelPtr> _channels;               // 当前服务对应的信道集合
    std::unordered_map<std::string, ChannelPtr> _hosts; // 主机地址与信道映射关系
};

// 总体的服务信道管理类
class ServiceManager {
public:
    using ptr = std::shared_ptr<ServiceManager>;

    // 构造函数
    ServiceManager() {}

    // 获取指定服务的节点信道
    ServiceChannel::ChannelPtr choose(const std::string &service_name) {
        std::unique_lock<std::mutex> lock(_mutex);
        auto sit = _services.find(service_name);
        if (sit == _services.end()) {
            LOG_ERROR("当前没有能够提供 {} 服务的节点！", service_name);
            return ServiceChannel::ChannelPtr();
        }
        return sit->second->choose();
    }

    // 声明需要关注哪些服务的上下线
    void declared(const std::string &service_name) {
        std::unique_lock<std::mutex> lock(_mutex);
        _follow_services.insert(service_name);
    }

    // 服务上线时调用的回调接口，将服务节点管理起来
    void onServiceOnline(const std::string &service_instance, const std::string &host) {
        std::string service_name = getServiceName(service_instance);
        ServiceChannel::ptr service;

        {
            std::unique_lock<std::mutex> lock(_mutex);
            auto fit = _follow_services.find(service_name);
            if (fit == _follow_services.end()) {
                LOG_DEBUG("{}-{} 服务上线了，但是当前并不关心！", service_name, host);
                return;
            }

            // 获取或创建服务的信道管理对象
            auto sit = _services.find(service_name);
            if (sit == _services.end()) {
                service = std::make_shared<ServiceChannel>(service_name);
                _services.insert(std::make_pair(service_name, service));
            } else {
                service = sit->second;
            }
        }

        if (!service) {
            LOG_ERROR("新增 {} 服务管理节点失败！", service_name);
            return;
        }

        // 添加新的服务节点信道
        service->append(host);
        LOG_DEBUG("{}-{} 服务上线新节点，进行添加管理！", service_name, host);
    }

    // 服务下线时调用的回调接口，从服务信道管理中删除指定节点信道
    void onServiceOffline(const std::string &service_instance, const std::string &host) {
        std::string service_name = getServiceName(service_instance);
        ServiceChannel::ptr service;

        {
            std::unique_lock<std::mutex> lock(_mutex);
            auto fit = _follow_services.find(service_name);
            if (fit == _follow_services.end()) {
                LOG_DEBUG("{}-{} 服务下线了，但是当前并不关心！", service_name, host);
                return;
            }

            // 查找服务信道管理对象
            auto sit = _services.find(service_name);
            if (sit == _services.end()) {
                LOG_WARN("删除{}服务节点时，没有找到管理对象", service_name);
                return;
            }
            service = sit->second;
        }

        // 从信道管理中移除服务节点
        service->remove(host);
        LOG_DEBUG("{}-{} 服务下线节点，进行删除管理！", service_name, host);
    }

private:
    // 从服务实例中提取服务名称
    std::string getServiceName(const std::string &service_instance) {
        auto pos = service_instance.find_last_of('/');
        if (pos == std::string::npos) return service_instance;
        return service_instance.substr(0, pos);
    }

private:
    std::mutex _mutex;                                     // 互斥锁，保护多线程环境下的成员变量
    std::unordered_set<std::string> _follow_services;      // 需要关注的服务名称集合
    std::unordered_map<std::string, ServiceChannel::ptr> _services; // 服务名称到信道管理对象的映射关系
};

} // namespace mag
